/* Copyright (c) 2024 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
 * Please refer to the License for details. You may not use this file except in compliance with the License.
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
 * See LICENSE in the root of the software repository for the full text of the License.
 * ===================================================================================================================*/

#ifndef AUTOFUSE_ASCIR_REGISTRY_H
#define AUTOFUSE_ASCIR_REGISTRY_H
#include <functional>
#include <sstream>
#include <string>
#include <unordered_map>
#include <vector>
namespace ge {
namespace ascir {
// TTODO: c++的类ABI兼容性不好，后面考虑换成C接口实现
struct AscIrAttrDef {
  std::string name;
  std::string asc_ir_type;
  std::string ge_ir_type;
};
struct AscIrDef {
  using CodeGenerator = void (*)(const AscIrDef &def, std::stringstream &ss);

  std::string file_path;
  int64_t line;
  std::string type;

  // 当前只有必选输入一种，没有其他类型，因此暂时简单处理，后续有复杂的optional后，defs的类型就不是string了
  std::vector<std::string> input_defs;
  std::vector<std::string> output_defs;
  std::vector<AscIrAttrDef> attr_defs;

  bool start_node;
  CodeGenerator infer_data_type_generator;
  CodeGenerator infer_view_generator;
};
inline void SameDataTypeFromInput(const AscIrDef &def, std::stringstream &ss, const char *input_name) {
  for (const auto &output_def : def.output_defs) {
    ss << "  op." << output_def << ".dtype = static_cast<ge::DataType>(" << input_name << "_in.dtype);" << std::endl;
  }
}
inline void SameDataTypeFromFirstInput(const AscIrDef &def, std::stringstream &ss) {
  if (!def.input_defs.empty()) {
    SameDataTypeFromInput(def, ss, def.input_defs[0].c_str());
  }
}
inline void SameDataTypeFromSecondInput(const AscIrDef &def, std::stringstream &ss) {
  if (def.input_defs.size() > 1U)
    SameDataTypeFromInput(def, ss, def.input_defs[1].c_str());
}
inline void SameViewFromInput(const AscIrDef &def, std::stringstream &ss, const char *input_name) {
  for (const auto &output_def : def.output_defs) {
    ss << "  op." << output_def << ".axis = " << input_name << "_in.axis;"
       << std::endl;
    ss << "  op." << output_def << ".repeats = " << input_name << "_in.repeats;"
       << std::endl;
    ss << "  op." << output_def << ".strides = " << input_name << "_in.strides;"
       << std::endl;
  }
}
inline void SameViewFromFirstInput(const AscIrDef &def, std::stringstream &ss) {
  if (!def.input_defs.empty()) {
    SameViewFromInput(def, ss, def.input_defs[0].c_str());
  }
}

class AscirRegistry {
 public:
  static AscirRegistry &GetInstance();
  void RegisterAscIr(const std::string &type, const AscIrDef &def);

  const std::unordered_map<std::string, AscIrDef> &GetAll() const;

 private:
  std::unordered_map<std::string, AscIrDef> types_to_ascir_;
};
}  // namespace ascir
}
#endif  // AUTOFUSE_ASCIR_REGISTRY_H
